# Author: qsj Time: 2023/02/25 Describe: (consult: Nadimi-Shahraki et.Enhanced whale optimization algorithm for
# medical feature selection: A COVID-19 case study[J].Computers in Biology and Medicine,2022.)

import math
import random

import numpy
from matplotlib import pyplot as plt

from testFunctions import f1, f5, f6


def poolMechanism(pool, X_best, X_worst, pool_size=100):
    # pool param initialization
    P_row = []

    pool_length = len(pool)
    free_space = pool_size - pool_length

    best_max = max(X_best)
    best_min = min(X_best)

    for i in range(len(X_worst)):
        B = [(0 if random.random() < 0.5 else 1) for j in range(len(X_best))]  # random binary vector

        # Random location near X_best
        X_brnd = [random.random() * (best_max - best_min) + best_min for j in range(len(X_best))]

        # cross
        Pi = [X_brnd[j] if B[j] == 1 else X_worst[i][j] for j in range(len(B))]
        P_row.append(Pi)

    # Exclude recurring elements
    P = []
    [P.append(pt) for pt in P_row if pt not in P]

    # Updata pool
    pool_row = []
    if free_space > 0:
        if free_space >= len(P):  # Enough space left for direct storage
            pool_row = pool + P
        else:  # Put in part, and the rest will randomly cover the history
            pool_row = pool + P[:free_space]
            rm = len(P) - free_space
            random_index_list = [random.randint(0, pool_length - 1) for i in range(rm)]
            for i, rix in enumerate(random_index_list):
                pool_row[rix] = P[free_space + i]
    else:
        pool_row = pool.copy()
        random_index_list = [random.randint(0, pool_length - 1) for i in range(len(P))]
        for i, rix in enumerate(random_index_list):
            pool_row[rix] = P[free_space + i]

    # Exclude recurring elements
    pool_res = []
    [pool_res.append(pt) for pt in pool_row if pt not in pool_res]

    return pool_res


def WhaleOptimizationAlgorithm_Enhanced(n=100, d=2, T=100, ub=None, lb=None, f=None):
    """
    Enhanced Whale Optimization Algorithm(E-WOA)
    origin:
    (1) paper: https://schlr.cnki.net/zn/Detail/index/GARJ2021_3/SJESA37290291496262A2B8A29ED514236E6
    (2) Matlab code: https://seyedalimirjalili.com/woa

    n: Number of population
    d: Individual position dimension(Solution dimension)
    T: Maximum Iterations
    ub: Upper boundary list
    lb: Lower boundary list
    f: Fitness function

    1. Please modify the calculation part of fitness function (y)
    2. Default is minimum optimization, change "best_score" and "y < best_score" to maximum optimization
    3.The return value is an iterative curve, which can be modified to the optimal solution if necessary
    """

    # param initialization(Find min value use inf)
    best_score = 1e10

    # Record best score of each iteration
    Convergence_curve = []

    # Position initialization
    Positions = [[random.uniform(lb[i], ub[i]) for i in range(d)] for j in range(n)]
    Positions = numpy.asarray(Positions)
    X_best = Positions[0].copy()

    # The rate of random select part
    ro = 0.2

    # Record fitness
    fm = [0 for i in range(n)]

    # pool initialization
    pool = []
    poolSize = int(1.5 * n)
    for i, x in enumerate(Positions):  fm[i] = f(Positions[i].copy())  # Calculate fitness
    # Sort the fitness from small to large and return the sorted index list(find min value use "reverse=False")
    sorted_index = sorted(range(len(fm)), key=lambda fmx: fm[fmx], reverse=False)
    # Set best position as leader(find min value use <)
    if fm[sorted_index[0]] < best_score:
        best_score = fm[sorted_index[0]]
        X_best = Positions[sorted_index[0]].copy()
    # Find the worst position
    worst_size = int(n - 0.3 * poolSize + 1)
    X_worst = [Positions[i] for i in sorted_index[-worst_size:]]
    # Pool Mechanism
    pool = poolMechanism(pool=pool, X_best=X_best.copy(), X_worst=X_worst, pool_size=poolSize)

    # Iteration start
    for t in range(T):

        # Use to record new position
        X = Positions.copy()

        # param
        a2 = -1 - t / T

        # Select a part of the population at random
        selected_index = random.sample([i for i in range(n)], int(ro * n))

        # Updata position of agent
        for i, x in enumerate(Positions):

            if i not in selected_index:

                p = random.random()

                A = 0.5 + 0.1 * math.tan(math.pi * (random.random() - 0.5))
                while A < 0:
                    A = 0.5 + 0.1 * math.tan(math.pi * (random.random() - 0.5))
                A = min(A, 1)

                C = 2 * random.random()

                if p < 0.5:
                    if abs(A) >= 0.5:  # Preferential selecting search strategy
                        r1 = random.randint(0, len(pool) - 1)
                        while r1 == i: r1 = random.randint(0, len(pool) - 1)
                        r2 = random.randint(0, len(pool) - 1)
                        while r2 == i or r2 == r1: r2 = random.randint(0, len(pool) - 1)
                        for j in range(d):
                            X[i][j] = Positions[i][j] + A * (C * pool[r1][j] - pool[r2][j])
                    else:  # Enriched encircling prey search strategy
                        for j in range(d):
                            r3 = random.randint(0, len(pool) - 1)
                            X[i][j] = X_best[j] - A * abs(C * X_best[j] - pool[r3][j])
                else:  # Spiral update position
                    for j in range(d):
                        b = 1
                        g = (a2 - 1) * random.random() + 1
                        Dl = abs(X_best[j] - Positions[i][j])
                        X[i][j] = Dl * math.exp(b * g) * math.cos(2 * math.pi * g) + X_best[j]
            else:  # Migrating search strategy
                X_rnd = [random.random() * (ub[k] - lb[k]) + lb[k] for k in range(d)]
                X_brnd = [random.random() * (max(X_best) - min(X_best)) + min(X_best)
                          for k in range(d)]
                X[i] = [X_rnd[k] - X_brnd[k] for k in range(d)]

        # Check bounds
        for i in range(n):
            for j in range(d):
                if X[i][j] > ub[j]:
                    X[i][j] = ub[j]
                elif X[i][j] < lb[j]:
                    X[i][j] = lb[j]

        # Updata best position
        X_worst = []
        for i in range(n):
            y = f(X[i].copy())  # Calculate fitness
            if y < best_score:
                best_score = y
                X_best = X[i].copy()

            # Compare individual fitness before and after iteration to update the position
            if y < fm[i]:
                # Put the poor position in the last iteration result into the X_worst
                X_worst.append(Positions[i].copy())
                # Updata position and fitness record
                Positions[i] = X[i].copy()
                fm[i] = y

        # Pool Mechanism
        pool = poolMechanism(pool=pool, X_best=X_best.copy(), X_worst=X_worst, pool_size=poolSize)

        # Record best position to updata iterative process curve
        Convergence_curve.append(best_score)

    # Return optimal solution
    return Convergence_curve


# # example
# population_numbers = 100
# dimension = 10
# iterations = int((10000 * dimension) / population_numbers)
#
# upper_bounds = [10 for i in range(dimension)]
# lower_bounds = [-10 for i in range(dimension)]
#
# xl = WhaleOptimizationAlgorithm_Enhanced(n=population_numbers, d=dimension, T=iterations,
#                                          ub=upper_bounds, lb=lower_bounds, f=f6)
#
# # Display the optimal result
# print("Optimization result: ", xl[-1])
#
# # Draw iterative change process curve
# plt.title('Iterative Process')
# plt.xlabel('iterations')
# plt.ylabel('best score')
# plt.plot([i for i in range(iterations)], xl)
# plt.show()



